#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   LSTMModel.py    
@Contact :   LJL959@QQ.com.com
@License :   (C)Copyright 2019-2020

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2020/4/7 9:49   LiuJiaoLong      1.0         None
@Description : 搭建LSTM网络
'''
import tensorflow as tf


def LSTMModel(num_of_word, input_data, output_data=None, run_size=128, num_layers=2, batch_size=128):
    end_points = {}
    """
    num_of_word:词的个数
    input_data:输入向量
    output_data:标签
    num_layers:隐藏层的层数
    batch_size:批处理数据大小
    end_points:存储训练得到的数据
    """
    # 构建LSTM功能的核心
    cell_fun = tf.contrib.rnn.BasicLSTMCell
    cell = cell_fun(run_size, state_is_tuple=True)
    # state_is_tuple=True的时候，state是元组，state=(c,h),为FALSE的时候，state=tf.concat(1,[c,h])
    # 构造完多层LSTM以后，使用zero_state即可对各种状态进行初始化
    # 使用了两层的LSTM网络，前一层的输出作为后一层的输入，用下述方法实现
    cell = tf.contrib.rnn.MultiRNNCell([cell] * num_layers, state_is_tuple=True)
    # 如果有标签（output_data）,则用zero_state初始化一个batch的cell状态，否则初始化为1
    if output_data is not None:
        initial_state = cell.zero_state(batch_size, tf.float32)
    else:
        initial_state = cell.zero_state(1, tf.float32)
    # 词向量的嵌入，构造一个num_of_word+1  X  run_size的矩阵，作为embedding容器，每个向量中的分量的值都在-1~1间随机分布
    embedding = tf.get_variable('embedding', initializer=tf.random_uniform([num_of_word+1, run_size], -1.0, 1.0))
    # 输入模型inputs，使用词典id转为向量
    inputs = tf.nn.embedding_lookup(embedding, input_data)
    # 只需使每个batch内数据长度相同即可，无需整个数据集内数据长度相同，所以使用dynamic_rnn，
    # 可以使不同迭代传入的batch是不同的长度，但同一次迭代一个batch内部的数据相同
    outputs, last_state = tf.nn.dynamic_rnn(cell, inputs, initial_state=initial_state)
    output = tf.reshape(outputs, [-1, run_size])
    weights = tf.Variable(tf.truncated_normal([run_size, num_of_word+1]))
    bias = tf.Variable(tf.zeros(shape=[num_of_word+1]))
    logits= tf.nn.bias_add(tf.matmul(output, weights), bias=bias)

    if output_data is not None:
        labels = tf.one_hot(tf.reshape(output_data, [-1]), depth=num_of_word+1)
        # 损失函数
        loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)
        total_loss = tf.reduce_mean(loss)   # 平均损失
        train_op = tf.train.AdamOptimizer(0.01).minimize(total_loss)  # 优化器
        tf.summary.scalar('loss', total_loss)
        end_points['initial_state'] = initial_state
        end_points['output'] = output
        end_points['train_op'] = train_op
        end_points['total_loss'] = total_loss
        end_points['loss'] = loss
        end_points['last_state'] = last_state
    else:
        prediction = tf.nn.softmax(logits)
        end_points['initial_state'] = initial_state
        end_points['last_state'] = last_state
        end_points['prediction'] = prediction

    return end_points
